from typing import Dict, List
from ..shard.placement_types import DeviceMesh,AbstractTensor,OperationData, OperationDataType
from .node_handler import NodeHandler
from ..generator import PlaceholderGenerator, StrategyGenerator,ReplicateGenerator
from ..StrategiesVector import StrategiesVector
from geesibling.core.types import Graph, Node
__all__ = ['PlaceholderHandler']


class PlaceholderHandler(NodeHandler):
    """
    A PlaceholderHandler which deals with the sharding strategies for Placeholder Node.
    """

    def __init__(self, node: Node, device_mesh: DeviceMesh, strategies_vector: StrategiesVector) -> None:
        super().__init__(node, device_mesh, strategies_vector)

    def get_strategy_generator(self) -> List[StrategyGenerator]:
        op_data_mapping = self.get_operation_data_mapping()
        generators = []
        if self.node.attrs['target'] == "USER_INPUT":
            #输入的策略，即数据并行
            # change for mistral arg11_1 need shard
            # generators.append(
            #     ReplicateGenerator(op_data_mapping, self.device_mesh))
            generators.append(
                PlaceholderGenerator(op_data_mapping, self.device_mesh))
        else:
            # 参数的切分策略
            generators.append(
                PlaceholderGenerator(op_data_mapping, self.device_mesh))
        return generators

    def get_operation_data_mapping(self) -> Dict[str, OperationData]:
        # use transposed shape for strategies
        # the strategies will be transformed back to its original shape in self.post_process
        # TODO 修改所有的data为TensorMeta
        output_abs_data = AbstractTensor(self.node.output_shape(0),self.node.output_type(0))
        physical_output = OperationData(name=str(self.node), data=output_abs_data,type=OperationDataType.OUTPUT)
        mapping = {"output": physical_output}
        return mapping
